# infer_config_dialog.py
import yaml
from PyQt5.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QFormLayout,
    QLabel, QLineEdit, QCheckBox, QPushButton, QWidget
)
from PyQt5.QtCore import Qt


class InferenceParamsDialog(QDialog):
    def __init__(self, yaml_path="infer.yaml", parent=None):
        super().__init__(parent)
        self.setWindowTitle("推理参数调整")
        self.resize(600, 400)
        self.setMinimumSize(600, 400)

        self.yaml_path = yaml_path
        self.params = {}
        self.inputs = {}

        # 参数说明
        self.param_desc = {
            "project": "推理结果输出的目录",
            "name": "推理实验名称",
            "show": "是否在窗口中显示检测结果",
            "save_json": "是否将检测结果保存为 JSON 文件",
            "save_txt": "是否将检测结果保存为 TXT 文件",
            "save_crop": "是否保存检测到的目标裁剪图",
            "show_conf": "是否在结果图上显示置信度",
            "show_boxes": "是否在结果图上显示检测框",
        }

        self._load_yaml()

        # ---------------- 布局 ----------------
        main_layout = QVBoxLayout()

        form_widget = QWidget()
        form_layout = QFormLayout()
        form_widget.setLayout(form_layout)

        for key in self.param_desc.keys():
            value = self.params.get(key, "")
            line_edit = QLineEdit(str(value))
            line_edit.setToolTip(self.param_desc[key])
            self.inputs[key] = line_edit

            label = QLabel(key)
            label.setToolTip(self.param_desc[key])
            form_layout.addRow(label, line_edit)

        main_layout.addWidget(form_widget)

        # 底部按钮
        btn_layout = QHBoxLayout()
        btn_layout.addStretch()
        self.btn_ok = QPushButton("确定")
        self.btn_cancel = QPushButton("取消")
        btn_layout.addWidget(self.btn_ok)
        btn_layout.addWidget(self.btn_cancel)
        main_layout.addLayout(btn_layout)

        self.setLayout(main_layout)

        # 事件绑定
        self.btn_ok.clicked.connect(self._on_ok)
        self.btn_cancel.clicked.connect(self.reject)

        # 非模态显示（这里可改模态显示）
        self.setWindowModality(Qt.ApplicationModal)

    def _load_yaml(self):
        """加载 infer.yaml 文件"""
        try:
            with open(self.yaml_path, "r", encoding="utf-8") as f:
                self.params = yaml.safe_load(f) or {}
        except Exception:
            self.params = {
                "project": "runs/predict",
                "name": "exp",
                "show": True,
                "save_json": False,
                "save_txt": False,
                "save_crop": False,
                "show_conf": True,
                "show_boxes": True,
            }

    def _on_ok(self):
        """保存修改到 YAML"""
        new_values = self.get_params()
        allowed_keys = list(self.param_desc.keys())

        try:
            with open(self.yaml_path, "r", encoding="utf-8") as f:
                lines = f.readlines()
        except Exception as e:
            print(f"读取 YAML 失败: {e}")
            self.accept()
            return

        for i, line in enumerate(lines):
            stripped = line.strip()
            if not stripped or stripped.startswith("#") or ":" not in stripped:
                continue
            key = stripped.split(":", 1)[0].strip()
            if key in allowed_keys and key in new_values:
                indent = line[:line.find(key)]
                comment = ""
                if "#" in line:
                    comment = line[line.find("#"):]
                lines[i] = f"{indent}{key}: {new_values[key]} {comment}".rstrip() + "\n"

        try:
            with open(self.yaml_path, "w", encoding="utf-8") as f:
                f.writelines(lines)
            print(f"推理参数已更新: {new_values}")
        except Exception as e:
            print(f"写入 YAML 失败: {e}")

        self.accept()

    def get_params(self):
        """返回用户修改后的参数"""
        new_params = {}
        for key, widget in self.inputs.items():
            text = widget.text()
            if text.lower() in ["true", "false"]:
                new_params[key] = text.lower() == "true"
            else:
                new_params[key] = text
        return new_params
